#if USE_VULKAN
#include "vkcontext.h"
#include "vkutils.h"
#include "vkblob.h"
#include "vkop.h"
#include "vkinfer/core/logger.h"
#include "vkinfer/core/utils.h"
#include "vkinfer/core/logger.h"

#include "vkinfer/ops/vk/vkop_add.h"
#include "vkinfer/ops/vk/vkop_relu.h"
#include "vkinfer/ops/vk/vkop_conv.h"
#include "vkinfer/ops/vk/vkop_linear.h"
#include "vkinfer/ops/vk/vkop_maxpool.h"
#include "vkinfer/ops/vk/vkop_averagepool.h"
#include "vkinfer/ops/vk/vkop_flatten.h"
#include "vkinfer/ops/vk/vkop_globalaveragepool.h"
#include "vkinfer/ops/vk/vkop_sigmoid.h"
#include "vkinfer/ops/vk/vkop_concat.h"

namespace vkinfer
{
    VkContext::VkContext()
    {
        create_instance();
        pick_physical_device();
        create_device();
        
        // create compute / transfer context.
        compute = std::make_shared<VkCompute>(device, compute_queue, compute_queue_index);
        transfer = std::make_shared<VkTransfer>(device, transfer_queue, transfer_queue_index);

        // create fence
        VkFenceCreateInfo fenceInfo{};
        fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
        fenceInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;
        VK_CHECK_RESULT(vkCreateFence(device, &fenceInfo, nullptr, &fence));
    }

    VkContext::~VkContext()
    {
        // destroy platform specific.
        vkDeviceWaitIdle(device);
        vkDestroyFence(device, fence, nullptr);
        compute.reset();
        transfer.reset();
        vkDestroyDevice(device, nullptr);
        vkDestroyInstance(instance, nullptr);
    }

    void VkContext::begin()
    {
        this->compute->begin_command_buffer();
    }

    void VkContext::end()
    {
        this->compute->end_command_buffer();
    }

    void VkContext::submit()
    {
        vkResetFences(device, 1, &fence);
        this->compute->submit(fence);
        // wait for computation finish.
        vkWaitForFences(device, 1, &fence, VK_TRUE, UINT64_MAX);
    }

    std::shared_ptr<Blob> VkContext::make_blob(const std::string& name, BlobTransferDirection direction)
    {
        // no shape information provided
        std::shared_ptr<VkBlob> blob = std::make_shared<VkBlob>(this->device, this->gpu, this->transfer, direction);
        blob->name = name;
        return blob;
    }

    std::shared_ptr<Blob> VkContext::make_blob(const std::string& name, uint32_t count, BlobTransferDirection direction)
    {
        std::vector<uint32_t> shape;
        shape.push_back(count);
        return this->make_blob(name, shape, direction);
    }

    std::shared_ptr<Blob> VkContext::make_blob(const std::string& name, const std::vector<uint32_t>& shape, BlobTransferDirection direction)
    {
        std::shared_ptr<VkBlob> blob = std::make_shared<VkBlob>(this->device, this->gpu, this->transfer, shape, direction);
        blob->name = name;
        return blob;
    }

    std::shared_ptr<Operator> VkContext::make_op(const std::string& name, const std::shared_ptr<OnnxNodeInfo>& node_info)
    {
        // for debug purpose
        OpType optype = op_string2type(node_info->op_type);

        if (optype == OpType::OP_TYPE_INVALID)
        {
            VK_ERROR("operator %s not supported\n", name.c_str());
            return nullptr;
        }

        std::shared_ptr<VkOp> op = nullptr;
        switch (optype)
        {
        case vkinfer::OpType::OP_TYPE_ADD:
            op = std::make_shared<VkOpAdd>(device, node_info);
            break;
        case vkinfer::OpType::OP_TYPE_RELU:
            op = std::make_shared<VkOpReLU>(device, node_info);
            break;
        case vkinfer::OpType::OP_TYPE_CONV2D:
            op = std::make_shared<VkOpConv>(device, node_info);
            break;
        case vkinfer::OpType::OP_TYPE_MAXPOOL:
            op = std::make_shared<VkOpMaxPool>(device, node_info);
            break;
        case vkinfer::OpType::OP_TYPE_AVERAGEPOOL:
            op = std::make_shared<VkOpAveragePool>(device, node_info);
            break;
        case vkinfer::OpType::OP_TYPE_LINEAR:
            op = std::make_shared<VkOpLinear>(device, node_info);
            break;
        case vkinfer::OpType::OP_TYPE_FLATTEN:
            op = std::make_shared<VkOpFlatten>(device, node_info);
            break;
        case vkinfer::OpType::OP_TYPE_GLOBALAVERAGEPOOL:
            op = std::make_shared<VkOpGlobalAveragePool>(device, node_info);
            break;
        case vkinfer::OpType::OP_TYPE_SIGMOID:
            op = std::make_shared<VkOpSigmoid>(device, node_info);
            break;
        case vkinfer::OpType::OP_TYPE_CONCAT:
            op = std::make_shared<VkOpConcat>(device, node_info);
            break;
        default:
            break;
        }
        if (op != nullptr)
            op->set_context(compute, transfer);
        return op;
    }

    // initialization
    void VkContext::create_instance()
    {
        std::vector<const char*> enabledLayers;


        uint32_t layerCount;
        vkEnumerateInstanceLayerProperties(&layerCount, nullptr);

        std::vector<VkLayerProperties> availableLayers(layerCount);
        vkEnumerateInstanceLayerProperties(&layerCount, availableLayers.data());

        for (const auto& layerProperties : availableLayers) {
            if (strcmp("VK_LAYER_KHRONOS_validation", layerProperties.layerName) == 0) {
                enabledLayers.push_back("VK_LAYER_KHRONOS_validation");
                break;
            }
        }
        // std::cout << "Validation layers enabled" << std::endl;
        // LOG_INFO("validation layer enabled\n");
        VkApplicationInfo applicationInfo = {};
        applicationInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
        applicationInfo.pApplicationName = "vulkan-compute";
        applicationInfo.applicationVersion = 0;
        applicationInfo.pEngineName = "vulkan-compute-engine";
        applicationInfo.engineVersion = 0;
        applicationInfo.apiVersion = VK_API_VERSION_1_2;

        VkInstanceCreateInfo createInfo = {};
        createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
        createInfo.flags = 0;
        createInfo.pApplicationInfo = &applicationInfo;

        createInfo.enabledLayerCount = (uint32_t)enabledLayers.size();
        createInfo.ppEnabledLayerNames = enabledLayers.data();

        createInfo.enabledExtensionCount = 0;
        createInfo.ppEnabledExtensionNames = NULL;

        VK_CHECK_RESULT(vkCreateInstance(&createInfo, NULL, &instance));
    }

    void VkContext::pick_physical_device()
    {
        uint32_t deviceCount;
        vkEnumeratePhysicalDevices(instance, &deviceCount, NULL);
        if (deviceCount == 0) 
        {
            std::cout << "Could not find a device with vulkan support" << std::endl;
            exit(0);
        }

        std::vector<VkPhysicalDevice> devices(deviceCount);
        vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());

#if 0
        for (const auto& dev : devices) 
        {
            if (is_device_suitable(dev)) 
            {
                gpu = dev;
                break;
            }
        }
#else
        // choose the better gpu.
        gpu = devices[devices.size() - 1];
#endif
        if (gpu == VK_NULL_HANDLE) 
        {
            std::cout << "Failed to find a suitable vulkan device" << std::endl;
            exit(-1);
        }
    }

    void VkContext::create_device()
    {
        std::vector<VkDeviceQueueCreateInfo> queueCreateInfos;

        VkDeviceQueueCreateInfo queue_ci = {};
        queue_ci.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
        QueueFamilyIndices indices = find_queue_families(gpu);

        // compute queue
        compute_queue_index = indices.compute_family.value();
        queue_ci.queueFamilyIndex = compute_queue_index;
        queue_ci.queueCount = 1;
        float queuePriorities = 1.0;
        queue_ci.pQueuePriorities = &queuePriorities;
        queueCreateInfos.push_back(queue_ci);

        // transfer queue
        if (transfer_queue_index != compute_queue_index)
        {
            transfer_queue_index = indices.transfer_family.value();
            queue_ci.queueFamilyIndex = transfer_queue_index;
            queue_ci.queueCount = 1;
            queuePriorities = 1.0;
            queue_ci.pQueuePriorities = &queuePriorities;
            queueCreateInfos.push_back(queue_ci);
        }

        VkDeviceCreateInfo deviceCreateInfo = {};
        VkPhysicalDeviceFeatures deviceFeatures = {};

        deviceCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
        deviceCreateInfo.enabledLayerCount = 0;
        deviceCreateInfo.ppEnabledLayerNames = NULL;
        deviceCreateInfo.pQueueCreateInfos = queueCreateInfos.data();
        deviceCreateInfo.queueCreateInfoCount = (uint32_t)queueCreateInfos.size();
        deviceCreateInfo.pEnabledFeatures = &deviceFeatures;

        std::vector<const char*> extensions;

#ifdef __APPLE__
        extensions.push_back("VK_KHR_portability_subset");
#endif

        deviceCreateInfo.ppEnabledExtensionNames = extensions.data();
        deviceCreateInfo.enabledExtensionCount = (uint32_t)extensions.size();

        VK_CHECK_RESULT(vkCreateDevice(gpu, &deviceCreateInfo, NULL, &device));

        vkGetDeviceQueue(device, compute_queue_index, 0, &compute_queue);
        vkGetDeviceQueue(device, transfer_queue_index, 0, &transfer_queue);
    }
}
#endif
